# 这里是多臂老虎机的案例代码
# 其余代码我直接写在了.ipynb文件里了
import numpy as np
import matplotlib.pyplot as plt

# 定义伯努利多臂老虎机
class BernoulliBandit:
    def __init__(self, K):
        self.probs = np.random.uniform(size=K)  # 随机生成各臂的获奖概率
        self.best_idx = np.argmax(self.probs)  # 最优臂索引
        self.best_prob = self.probs[self.best_idx]  # 最优臂概率
        self.K = K  # 臂的数量

    def step(self, k):
        return 1 if np.random.rand() < self.probs[k] else 0  # 根据概率返回奖励

# 定义算法基类
class Solver:
    def __init__(self, bandit):
        self.bandit = bandit
        self.counts = np.zeros(self.bandit.K)  # 各臂尝试次数
        self.regret = 0.0  # 累积懊悔
        self.actions = []  # 动作记录
        self.regrets = []  # 懊悔记录

    def update_regret(self, k):
        self.regret += self.bandit.best_prob - self.bandit.probs[k]
        self.regrets.append(self.regret)

    def run_one_step(self):
        raise NotImplementedError

    def run(self, num_steps):
        for _ in range(num_steps):
            k = self.run_one_step()
            self.counts[k] += 1
            self.actions.append(k)
            self.update_regret(k)

# 实现epsilon-贪婪算法
class EpsilonGreedy(Solver):
    def __init__(self, bandit, epsilon=0.01, init_prob=1.0):
        super().__init__(bandit)
        self.epsilon = epsilon
        self.estimates = np.array([init_prob] * self.bandit.K)  # 初始估计值

    def run_one_step(self):
        if np.random.random() < self.epsilon:  # 探索
            k = np.random.randint(0, self.bandit.K)
        else:  # 利用
            k = np.argmax(self.estimates)
        
        r = self.bandit.step(k)  # 获取奖励
        self.estimates[k] += 1.0 / (self.counts[k] + 1) * (r - self.estimates[k])  # 更新估计值
        return k

# 绘图函数
def plot_results(solvers, solver_names):
    for idx, solver in enumerate(solvers):
        plt.plot(range(len(solver.regrets)), solver.regrets, label=solver_names[idx])
    
    plt.xlabel('Time steps')
    plt.ylabel('Cumulative regrets')
    plt.title(f'{solvers[0].bandit.K}-armed bandit')
    plt.legend()
    plt.show()

# 实验设置
np.random.seed(1)
K = 10
bandit = BernoulliBandit(K)
print(f"随机生成了一个{K}臂伯努利老虎机")
print(f"最优臂为{bandit.best_idx}号，获奖概率为{bandit.best_prob:.4f}")

# 运行算法
epsilon_greedy = EpsilonGreedy(bandit, epsilon=0.01)
epsilon_greedy.run(5000)
print(f'epsilon-贪婪算法的累积懊悔为: {epsilon_greedy.regret:.4f}')

# 绘制结果
plot_results([epsilon_greedy], ["EpsilonGreedy"])